use core::cmp::Ordering;

use alloc::collections::binary_heap::BinaryHeap;
use critical_section::with;
use utils::MutCell;

use crate::{osWaitForever, os_kernel_get_tick_count, Instant};

use super::EventFlags;

struct DelayItem<T> {
    value: T,
    insert: u32,
    delay: u32,
}

const FLAG_NOTIFY: u32 = 0x01;

impl<T> Eq for DelayItem<T> {}

impl<T> PartialEq for DelayItem<T> {
    fn eq(&self, other: &Self) -> bool {
        (self.delay == other.delay) && (self.insert == other.insert)
    }
}

impl<T> PartialOrd for DelayItem<T> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        let n0 = (self.delay as u64) + (self.insert as u64);
        let n1 = (other.delay as u64) + (other.insert as u64);
        n1.partial_cmp(&n0)
    }
}

impl<T> Ord for DelayItem<T> {
    fn cmp(&self, other: &Self) -> Ordering {
        let n0 = (self.delay as u64) + (self.insert as u64);
        let n1 = (other.delay as u64) + (other.insert as u64);
        n1.cmp(&n0)
    }
}

pub struct DelayQueue<T> {
    queue: MutCell<BinaryHeap<DelayItem<T>>>,
    event_flags: EventFlags,
}

impl<T: Send + 'static> DelayQueue<T> {
    pub fn new() -> Self {
        Self {
            queue: MutCell::const_new(BinaryHeap::new()),
            event_flags: EventFlags::new("delay queue"),
        }
    }

    pub fn post(&self, delay: u32, value: T) {
        let item = DelayItem {
            value,
            insert: os_kernel_get_tick_count(),
            delay,
        };
        with(|_| {
            unsafe { self.queue.as_mut() }.push(item);
        });
        _ = self.event_flags.set_flags(FLAG_NOTIFY);
    }

    pub fn take(&self) -> T {
        loop {
            let current = os_kernel_get_tick_count();
            let item = with(|_| {
                let queue = unsafe { self.queue.as_mut() };
                if let Some(item) = queue.peek() {
                    if current.wrapping_sub(item.insert) >= item.delay {
                        return queue.pop()
                    } 
                } 
                None
            });
            if let Some(item) = item {
                return item.value;
            }
            _ = self.event_flags.wait_any_flags(FLAG_NOTIFY, osWaitForever);
        }
    }

    pub fn poll(&self, timeout: u32) -> Option<T> {
        let mut time_count = 0;

        loop {
            let current = os_kernel_get_tick_count();
            let item = with(|_| {
                let queue = unsafe { self.queue.as_mut() };
                if let Some(item) = queue.peek() {
                    if current.wrapping_sub(item.insert) >= item.delay {
                        return queue.pop()
                    } 
                } 
                None
            });
            if let Some(item) = item {
                return Some(item.value);
            }
            if time_count >= timeout {
                return None;
            }
            let instant = Instant::new();
            _ = self.event_flags.wait_any_flags(FLAG_NOTIFY, timeout - time_count);
            time_count += instant.elapsed();
        }
    }
}
